import numpy as np
from pymoo.algorithms.soo.nonconvex.ga import GA
from pymoo.algorithms.soo.nonconvex.de import DE
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.core.problem import Problem
from pymoo.operators.crossover.pntx import TwoPointCrossover  # GA专用交叉算子
from pymoo.operators.mutation.pm import PolynomialMutation    # GA专用变异算子
from pymoo.operators.selection.tournament import TournamentSelection  # 选择算子
from pymoo.optimize import minimize
import pandas as pd
import time
# =======================海南测试=====================
# ====================== 基础参数设置 ======================
T = 120  # 3天×24小时×2时段/小时
dt = 3600  # 时段长度（秒）

def storage_hl_func(h_array: np.ndarray) -> np.ndarray:

    """
    水位转库容函数
    向量化计算存储函数值（支持ndarray输入）
    参数:
    h_array (ndarray): 待计算的容积数组（单位：万m³）
    返回:
    ndarray: 计算结果数组（超出区间的值为np.nan）
    """
    # 转换为浮点数组避免整数运算问题
    h = h_array.astype(np.float64)

    # 定义所有区间条件（按顺序）
    cond1 = np.logical_and(h >= -np.inf, h <= 92)          # 区间 [90, 92]
    cond2 = np.logical_and(h > 92, h <= 107)       # 区间 (92.01, 107]
    cond3 = np.logical_and(h > 107, h <= 116)      # 区间 (107.01, 116]
    cond4 = np.logical_and(h > 116, h <= 126)      # 区间 (116.01, 126]
    cond5 = np.logical_and(h > 126, h <= 136)      # 区间 (126.01, 136]
    cond6 = np.logical_and(h > 136, h <= 146)      # 区间 (136.01, 146]
    cond7 = np.logical_and(h > 146, h <= 156)      # 区间 (146.01, 156]
    cond8 = np.logical_and(h > 156, h <= 166)      # 区间 (156.01, 166]
    cond9 = np.logical_and(h > 166, h <= 169)      # 区间 (166.01, 169]
    cond10 = np.logical_and(h > 169, h <= np.inf)     # 区间 (169.01, 174]

    # 计算各区间的多项式值（顺序与条件一一对应）
    poly1 = 0.0021 * (h ** 2) + 9.2788 * h - 851.93
    poly2 = 3.8886 * (h ** 2) - 712.61 * h + 32666
    poly3 = 5.9448 * (h ** 2) - 1152.8 * h + 56235
    poly4 = 9.053 * (h ** 2) - 1871.5 * h + 97814
    poly5 = 12.752 * (h ** 2) - 2794.2 * h + 155382
    poly6 = 16.589 * (h ** 2) - 3837.7 * h + 226371
    poly7 = 25.892 * (h ** 2) - 6546.9 * h + 423702
    poly8 = 39.559 * (h ** 2) - 10809 * h + 756094
    poly9 = 2545.2 * h - 370375
    poly10 = 3161.3 * h - 475013

    # 根据条件选择结果（超出区间返回np.nan）
    result = np.select(
        condlist=[cond1, cond2, cond3, cond4, cond5, cond6, cond7, cond8, cond9, cond10],
        choicelist=[poly1, poly2, poly3, poly4, poly5, poly6, poly7, poly8, poly9, poly10],
        default=np.nan
    )

    return result
def storage_hl_inh_func(v_array: np.ndarray) -> np.ndarray:
    """
    库容转水位
    向量化计算存储函数值（支持ndarray输入）

    参数:
    v_array (ndarray): 待计算的容积数组（单位：万m³）

    返回:
    ndarray: 计算结果数组（超出区间或间隙值为np.nan）
    """
    # 转换为双精度浮点数组避免整数运算问题
    v = v_array.astype(np.float64)

    # 定义各区间的条件（按顺序）
    cond1 = np.logical_and(v >= -np.inf, v <= 48.3)          # 区间 [0, 48.3]
    cond2 = np.logical_and(v > 48.3, v <= 951.1)       # 区间 (48.8, 951.1]
    cond3 = np.logical_and(v > 951.1, v <= 2519.1)     # 区间 (952.4, 2519.1]
    cond4 = np.logical_and(v > 2519.1, v <= 5750.4)    # 区间 (2521.7, 5750.4]
    cond5 = np.logical_and(v > 5750.4, v <= 11259.3)     # 区间 (5755, 11259.3]
    cond6 = np.logical_and(v > 11259.3, v <= 19734.4)  # 区间 (11266.6, 19734.4]
    cond7 = np.logical_and(v > 19734.4, v <= 32563.3)  # 区间 (19745.4, 32563.3]
    cond8 = np.logical_and(v > 32563.3, v <= 52079.5)    # 区间 (32580, 52079.5]
    cond9 = np.logical_and(v > 52079.5, v <= 59661.3)  # 区间 (52105.7, 59661.3]
    cond10 = np.logical_and(v > 59661.3, v <= np.inf)   # 区间 (59685, 78135.9]

    # 计算各区间的多项式值（顺序与条件严格对应）
    poly1 = 7e-8 * (v ** 2) + 0.1035 * v + 90
    poly2 = 1e-8 * (v ** 3) - 2e-5 * (v ** 2) + 0.0281 * v + 93.617
    poly3 = -1e-6 * (v ** 2) + 0.0095 * v + 99.08
    poly4 = -3e-7 * (v ** 2) + 0.0053 * v + 104.32
    poly5 = -8e-8 * (v ** 2) + 0.0031 * v + 110.63
    poly6 = -3e-8 * (v ** 2) + 0.002 * v + 116.56
    poly7 = -1e-8 * (v ** 2) + 0.0014 * v + 122.62
    poly8 = -5e-9 * (v ** 2) + 0.001 * v + 130.01
    poly9 = 0.0004 * v + 145.53
    poly10 = 0.0003 * v + 150.27

    # 根据条件选择结果（区间间隙/超出范围返回np.nan）
    result = np.select(
        condlist=[cond1, cond2, cond3, cond4, cond5, cond6, cond7, cond8, cond9, cond10],
        choicelist=[poly1, poly2, poly3, poly4, poly5, poly6, poly7, poly8, poly9, poly10],
        default=np.nan
    )

    return result

def storage_nll_func(h_array: np.ndarray) -> np.ndarray:

    """
    水位转库容函数
    向量化计算存储函数值（支持ndarray输入）
    参数:
    h_array (ndarray): 待计算的容积数组（单位：万m³）
    返回:
    ndarray: 计算结果数组（超出区间的值为np.nan）
    """
    # 转换为双精度浮点数组避免整数运算问题
    h = h_array.astype(np.float64)

    # 定义各区间的条件（按顺序）
    cond1 = np.logical_and(h >= 70, h <= 74.9)  # 区间 [70, 74.9]
    cond2 = np.logical_and(h >= 75, h <= 79.9)  # 区间 [75, 79.9]
    cond3 = np.logical_and(h >= 80, h <= 84.9)  # 区间 [80, 84.9]
    cond4 = np.logical_and(h >= 85, h <= 89.9)  # 区间 [85, 89.9]
    cond5 = np.logical_and(h >= 90, h <= 94.9)  # 区间 [90, 94.9]
    cond6 = np.logical_and(h >= 95, h <= 99.9)  # 区间 [95, 99.9]
    cond7 = np.logical_and(h >= 100, h <= 104.9)  # 区间 [100, 104.9]
    cond8 = np.logical_and(h >= 105, h <= 109.9)  # 区间 [105, 109.9]

    # 计算各区间的多项式值（顺序与条件严格对应）
    poly1 = 13.197 * (h ** 2) - 1403.9 * h + 38897
    poly2 = 23.564 * (h ** 2) - 2966.8 * h + 97801
    poly3 = 29.042 * (h ** 2) - 3859.5 * h + 134140
    poly4 = 42.715 * (h ** 2) - 6329 * h + 245243
    poly5 = 46.552 * (h ** 2) - 7071 * h + 280960
    poly6 = 39.447 * (h ** 2) - 5654.1 * h + 210441
    poly7 = 62.235 * (h ** 2) - 10089 * h + 426195
    poly8 = 47.831 * (h ** 2) - 6939.1 * h + 254213

    # 根据条件选择结果（超出区间返回np.nan）
    result = np.select(
        condlist=[cond1, cond2, cond3, cond4, cond5, cond6, cond7, cond8],
        choicelist=[poly1, poly2, poly3, poly4, poly5, poly6, poly7, poly8],
        default=np.nan
    )

    return result
def storage_nll_inh_func(v_array: np.ndarray)-> np.ndarray:
    """
    库容转水位
    向量化计算存储函数值（支持ndarray输入）

    参数:
    v_array (ndarray): 待计算的容积数组（单位：万m³）

    返回:
    ndarray: 计算结果数组（超出区间或间隙值为np.nan）
    """
    # 转换为双精度浮点数组
    v = v_array.astype(np.float64)

    # 定义各区间条件（严格闭区间 [a, b]）
    cond1 = np.logical_and(v >= -np.inf, v <= 7785)  # 区间 [5289, 7785]
    cond2 = np.logical_and(v > 7785, v <= 11179)  # 区间 [7843, 11179]
    cond3 = np.logical_and(v > 11179, v <= 15775)  # 区间 [11255, 15775]
    cond4 = np.logical_and(v > 15775, v <= 21492)  # 区间 [15875, 21492]
    cond5 = np.logical_and(v > 21492, v <= 29153)  # 区间 [21625, 29153]
    cond6 = np.logical_and(v > 29153, v <= 39321)  # 区间 [29325, 39321]
    cond7 = np.logical_and(v > 39321, v <= 52677)  # 区间 [39562, 52677]
    cond8 = np.logical_and(v > 52677, v <= np.inf)  # 区间 [52975, 69330]

    # 定义各区间对应的多项式计算（顺序与条件严格对应）
    poly1 = -1e-7 * (v ** 2) + 0.0033 * v + 55.51
    poly2 = -7e-8 * (v ** 2) + 0.0029 * v + 57.139
    poly3 = -4e-8 * (v ** 2) + 0.002 * v + 61.546
    poly4 = -3e-8 * (v ** 2) + 0.0019 * v + 61.435
    poly5 = -1e-8 * (v ** 2) + 0.0013 * v + 67.909
    poly6 = -5e-9 * (v ** 2) + 0.0008 * v + 75.232
    poly7 = -3e-9 * (v ** 2) + 0.0007 * v + 78.309
    poly8 = -1e-9 * (v ** 2) + 0.0005 * v + 84.459

    # 向量化选择结果（区间间隙和超出范围自动填充np.nan）
    result = np.select(
        condlist=[cond1, cond2, cond3, cond4, cond5, cond6, cond7, cond8],
        choicelist=[poly1, poly2, poly3, poly4, poly5, poly6, poly7, poly8],
        default=np.nan
    )

    return result
# 水库参数（保持不变）
reservoir_params = {
    'A': {
        'Q_min': 0,  # 最小出库流量(m³/s)
        'Q_max': 7000,  # 最大出库流量(m³/s)
        'H_initial': 168.28,
        'V_initial': 57961.2    # 红岭初始库容, m³
    },
    'B': {
        'Q_min': 0,
        'Q_max': 8000,
        'H_initial': 103.59,
        'V_initial': 48848    # 牛路岭初始库容, m³
    }
}
# 断面参数
section_params = {
    'gate_JJ': {
        'H_initial': 2.5,
        'Max_flow': 16000,
        'Q_to_H_func': lambda Q: (Q + 12955)/ 1552.8
    }

}
# 水闸流达时间（分钟）→ 转换为时段数（半小时为一时段）
travel_time = {
    '牛路岭到嘉积': 7,  # 牛路岭到嘉积
    '红岭到嘉积': 8,  # 红岭到嘉积
}
# 转换为时段数（向上取整）
travel_time_period = {k: int(np.ceil(v / 1)) for k, v in travel_time.items()}


# 水闸列表（按下游顺序）
# downstream_gates = ['布吉河口', '鹿丹村', '上步码头', '深圳河口']
def load_inflow_data(excel_path):
    try:
        # 从Excel读取数据，假设Excel文件包含'A_inflow'和'B_inflow'两列
        df = pd.read_excel(excel_path)

        # 检查数据行数是否为时段数
        if len(df) != T:
            raise ValueError(f"数据行数应为{T}行，当前读取到{len(df)}行")

        # 检查是否包含所需列
        required_columns = ['时间','牛路入库', '红岭入库']
        if not set(required_columns).issubset(df.columns):
            missing = [col for col in required_columns if col not in df.columns]
            raise KeyError(f"Excel文件缺少必要列: {missing}")

        return df[required_columns]

    except FileNotFoundError:
        print(f"错误：未找到文件 {excel_path}")
        return None
    except Exception as e:
        print(f"读取数据时发生错误: {str(e)}")
        return None


# 遇洪水过程 5天
excel_file = "D:\code\Dispatch-SZW\src\data-hainan\两水库输入.xlsx"
inflow_data = load_inflow_data(excel_file)
A_lateral_inflow = np.zeros(T)  # A水库旁侧入流（来水）
B_lateral_inflow = np.zeros(T)  # B旁侧入流（来水）
if inflow_data is not None:
    print("成功读取入库流量数据：")
    print(inflow_data.head())  # 打印前5行验证


def compute_flow(upstream_flow, delay, total_time):
    """
    通用流量计算函数
    :param upstream_flow: 上游流量数组
    :param delay: 延迟时段
    :param total_time: 总时间步长T
    :return: 当前节点流量数组
    """
    return np.array([
        upstream_flow[t - delay] if t >= delay else 0
        for t in range(total_time)
    ])


# ====================== 模型 ======================
class ReservoirOptimizationProblem(Problem):
    def __init__(self):
        super().__init__(
            n_var=2 * T,  # 决策变量：A水库T时段出库 + B水库T时段出库
            n_obj=1,  # 单目标优化
            n_constr=14 * T,  # 约束
            xl=np.array([reservoir_params['A']['Q_min']] * T + [reservoir_params['B']['Q_min']] * T),  # 出库下限
            xu=np.array([reservoir_params['A']['Q_max']] * T + [reservoir_params['B']['Q_max']] * T)  # 出库上限
        )
        self.flow_processes = []  # 保存水库、断面的流量过程

    def _evaluate(self, X, out, *args, **kwargs):
        f = []  # 目标函数值（洪峰流量）
        g = []  # 约束条件

        for x in X:
            # 分解决策变量
            Q_A = x[:T]  # 红岭水库A各时段出库流量(m³/s)
            Q_B = x[T:]  # 牛路岭水库B各时段出库流量(m³/s)

            # ---------------------- 水库调蓄计算 ----------------------
            # 红岭水库A水量平衡
            V_A = np.zeros(T + 1)  # 库容序列 万m³
            H_A = np.zeros(T + 1)  # 红岭水库水位 m
            V_A[0] = reservoir_params['A']['V_initial']  # 红岭水库初始流量
            H_A[0] = reservoir_params['A']['H_initial']  # 红岭水库初始水位

            # 牛路岭水库B水量平衡
            V_B = np.zeros(T + 1)  # 红岭库容序列（万m³）
            H_B = np.zeros(T + 1)  # 红岭水位 m
            V_B[0] = reservoir_params['B']['V_initial']  # 牛路岭水库初始水位
            H_B[0] = reservoir_params['B']['H_initial']  # 牛路岭水库初始水位
            for t in range(T):
                inflow = inflow_data['红岭入库'].iloc[t]  # 红岭入库流量(m³/s)
                inflow1 = inflow_data['牛路入库'].iloc[t]  # 牛路岭入库流量(m³/s)
                V_A[t + 1] = V_A[t] + (inflow - Q_A[t]) * dt / 10000  # 转换为万m³
                H_A[t + 1] = storage_hl_inh_func(V_A[t + 1])

                # if V_B[t] < 311:
                #     Q_B[t] = 0
                #     V_B[t + 1] = V_B[t] + (inflow1 - Q_B[t]) * dt / 10000
                #     H_B[t + 1] = storage_hl_func(V_B[t + 1])
                # else:
                V_B[t + 1] = V_B[t] + (inflow1 - Q_B[t]) * dt / 10000
                H_B[t + 1] = storage_nll_inh_func(V_B[t + 1])

            # ---------------------- 下游水闸流量计算 ----------------------
            # 1. 计算布罗湖站流量（A和B的汇流）
            # A到罗湖站的延迟时段路径：A→三岔河口→文锦渡→罗湖站
            delay_dict = {
                '牛路岭到嘉积': (travel_time_period['牛路岭到嘉积']),
                '红岭到嘉积': travel_time_period['红岭到嘉积'],
            }

            # 1. A红岭到嘉积，
            Q_A_to_JJ = compute_flow(Q_A, delay_dict['牛路岭到嘉积'], T)
            # 2.牛路岭水库到布吉河口的流量
            Q_B_to_JJ = compute_flow(Q_B, delay_dict['红岭到嘉积'], T)
            Q_JJ = np.zeros(T + 1)  # 扩展为T+1个时段（含初始时刻）
            Q_JJ[0] = 0  # 扩展为T+1个时段（含初始时刻）
            for t in range(T):
                Q_JJ[t+1] = Q_A_to_JJ[t] + Q_B_to_JJ[t]
            # 嘉积水位
            H_JJ = np.zeros(T + 1)  # 扩展为T+1个时段（含初始时刻）
            H_JJ[0] = section_params['gate_JJ']['Q_to_H_func'](0)  # 初始水位（t=0）
            for t in range(T):
                # t时段的出库流量经汇流后影响t+1时段的H_JJ（与水库调蓄逻辑一致）
                H_JJ[t + 1] = section_params['gate_JJ']['Q_to_H_func'](Q_JJ[t+1])
            # for t in range(T):
            #     if t == 0:
            #         H_JJ[t] = section_params['gate_JJ']['H_initial']
            #     elif Q_BJHK[t] == 0:
            #         H_BJHK[t] = H_BJHK[t - 1]
            #     else:
            #         H_BJHK[t] = section_params['gate_BJHK']['Q_to_H_func'](Q_BJHK[t])
            # ---------------------- 目标函数：下游最大洪峰流量 ----------------------
            # 嘉积大坝的洪峰流量（过程最大值）

            peak_JJ = np.max(Q_JJ)
            # 下游整体最大洪峰
            total_peak_q = np.max([peak_JJ])
            total_peak_h = np.max([H_JJ])


            # total_v = 0
            # for t in range(T):
            #     total_v += (V_B[t+1] - V_B[t]) ** 2

            # 水位波动最小
            total_h = 0
            for t in range(T):
                total_h = H_JJ[t+1] - H_JJ[t]
            # 出库流量最小
            total_q = 0

            # 整体流量波动最小
            fluctuation_Q = 0
            for t in range(T-1):
                fluctuation_Q += Q_A[t+1] - Q_A[t] + Q_B[t+1] - Q_B[t]

            total_v = 0
            for t in range(T):
                total_v += (
                            (V_A[t+1] - reservoir_params['A']['V_initial']) ** 2
                            + (V_B[t+1] - reservoir_params['B']['V_initial']) ** 2
                            )

            total_b_v = 0
            for t in range(T):
                total_b_v += V_B[t + 1]
            f.append(total_h)

            # ---------------------- 约束条件 ----------------------
            # 1. 红岭水库时段末库容非负（V(t+1) ≥ 0）
            constr_vA = -V_A[1:]  # -V ≤ 0 → V ≥ 0

            # 2. 牛路岭水库时段末库容非负
            constr_vB = -V_B[1:]
            # 3. 红岭水库水位不超设计上限
            constr_zA = H_A[1:] - 171  # Z - Z_max ≤ 0
            constr_zA1 = 160 - H_A[1:]
            # 4. 牛路岭水库水位不超设计上限
            constr_zB = H_B[1:] - 114
            # 5. 牛路岭水库水位不超设计下限
            constr_zB1 = 99.85 - H_B[1:]
            # 嘉积大坝水位上下限
            constr_zC = H_JJ[1:] - 20
            constr_zC1 = -H_JJ[1:]

            constr_qC = -Q_JJ[1:]
            constr_qC1 = Q_JJ[1:] - 16000

            # 水位变幅约束
            # A水库水位变幅约束：|H_A[t+1] - H_A[t]| ≤ 1
            dH_A = H_A[1:] - H_A[:T]  # 时段t到t+1的水位变化（T个元素）
            constr_dH_A_max = dH_A - 2  # H_A[t+1] - H_A[t] ≤ 1
            constr_dH_A_min = -dH_A - 2  # H_A[t] - H_A[t+1] ≤ 1（等价于H_A[t+1] ≥ H_A[t] - 1）

            # B水库水位变幅约束：|H_B[t+1] - H_B[t]| ≤ 1
            dH_B = H_B[1:] - H_B[:T]
            constr_dH_B_max = dH_B - 2
            constr_dH_B_min = -dH_B - 2

            # C水库水位变幅约束：|H_C[t+1] - H_C[t]| ≤ 1
            dH_C = H_JJ[1:] - H_JJ[:T]
            constr_dH_C_max = dH_C - 0.7
            constr_dH_C_min = -dH_C - 0.7
            # C水库流量变幅约束：|H_C[t+1] - H_C[t]| ≤ 1
            dQ_C = Q_JJ[1:] - Q_JJ[:T]
            constr_dQ_C_max = dQ_C - 4000
            constr_dQ_C_min = -dQ_C - 4000

            # 合并所有约束（g ≤ 0 表示可行）
            g.append(np.concatenate([
                                     # constr_vA,
                                     # constr_vB,
                                     # constr_vB1,
                                     constr_zA,
                                     constr_zA1,
                                     constr_zB,
                                     constr_zB1,
                                     constr_zC,
                                     constr_zC1,
                                     constr_qC,
                                     constr_qC1,
                                     constr_dH_A_max,
                                     constr_dH_A_min,
                                     constr_dH_B_max,
                                     constr_dH_B_min,
                                     constr_dH_C_max,
                                     constr_dH_C_min,
                                     ])
                     )
            # 保存当前解的流量过程
            self.flow_processes.append({
                '红岭水库出库流量': Q_A,
                '牛路岭水库出库流量': Q_B,
                '红岭水库水位': H_A[1:],
                '红岭水库库容': V_A[1:],
                '牛路岭水库水位': H_B[1:],
                '牛路岭水库库容': V_B[1:],
                '嘉积流量': Q_JJ,
                '嘉积水位': H_JJ,
            })
        out["F"] = np.array(f).reshape(-1, 1)  # 目标函数值
        out["G"] = np.array(g)  # 约束条件


# ====================== 打印优化过程 ======================
def optimization_callback(algorithm):
    """回调函数：输出每代优化进度"""
    gen = algorithm.n_gen
    current_fitness = np.min(algorithm.pop.get("F"))
    print(f"第 {gen:3d} 代 | 当前适应度: {current_fitness:.2f}")


if __name__ == "__main__":
    start_time = time.perf_counter()  # 记录开始时间（高精度计时器）
    # 初始化优化问题
    problem = ReservoirOptimizationProblem()

    # 配置差分进化算法
    # algorithm = GA(
    #     pop_size=100,  # 增大种群大小（原50→100）提升多样性搜索能力
    #     selection=TournamentSelection(func_comp=lambda a, b: a.get("F") < b.get("F")),  # 锦标赛选择
    #     crossover=TwoPointCrossover(prob=0.8),  # 两点交叉（概率0.8）
    #     mutation=PolynomialMutation(prob=0.1, eta=20),  # 多项式变异（概率0.1，eta控制变异步长）
    #     eliminate_duplicates=True,  # 消除重复个体提升效率
    #     callback_data=[],  # 用于存储回调数据的列表
    # )
    algorithm = DE(
        pop_size=50,  # 种群大小
        variant="DE/rand/1/bin",  # 算法变体
        CR=0.7,  # 交叉概率
        F=0.5,  # 缩放因子
        dither="vector",  # 增加扰动提高多样性
        jitter=False
    )

    # 运行优化
    res = minimize(
        problem,
        algorithm,
        termination=('n_gen', 3000),  # 迭代次数
        callback=optimization_callback,  # 显示优化过程
        seed=123,  # 随机种子
        verbose=False  # 关闭默认输出
    )

    print(f"最终种群目标值: {res.F}")
    # 输出结果
    # print(f"\n优化完成,最终最小水位违反值: {res.F[0]:.2f} m")

    # 保存最优出库方案
    optimal_Q_A = res.X[:T]  # 红岭水库A最优出库序列
    optimal_Q_B = res.X[T:]  # 牛路岭水库B最优出库序列

    # ---------------------- 水库调蓄计算 ----------------------
    # 红岭水库A水量平衡
    V_A = np.zeros(T + 1)  # 库容序列 万m³
    H_A = np.zeros(T + 1)  # 红岭水库水位 m
    V_A[0] = reservoir_params['A']['V_initial']  # 红岭水库初始流量
    H_A[0] = reservoir_params['A']['H_initial']  # 红岭水库初始水位

    # 牛路岭水库B水量平衡
    V_B = np.zeros(T + 1)  # 红岭库容序列（万m³）
    H_B = np.zeros(T + 1)  # 红岭水位 m
    V_B[0] = reservoir_params['B']['V_initial']  # 牛路岭水库初始水位
    H_B[0] = reservoir_params['B']['H_initial']  # 牛路岭水库初始水位
    for t in range(T):
        inflow = inflow_data['红岭入库'].iloc[t]  # 红岭入库流量(m³/s)
        inflow1 = inflow_data['牛路入库'].iloc[t]  # 牛路岭入库流量(m³/s)
        V_A[t + 1] = V_A[t] + (inflow - optimal_Q_A[t]) * dt / 10000  # 转换为万m³
        H_A[t + 1] = storage_hl_inh_func(V_A[t + 1])

        # if V_B[t] < 311:
        #     Q_B[t] = 0
        #     V_B[t + 1] = V_B[t] + (inflow1 - Q_B[t]) * dt / 10000
        #     H_B[t + 1] = storage_hl_func(V_B[t + 1])
        # else:
        V_B[t + 1] = V_B[t] + (inflow1 - optimal_Q_B[t]) * dt / 10000
        H_B[t + 1] = storage_nll_inh_func(V_B[t + 1])

    # ---------------------- 下游水闸流量计算 ----------------------
    # 1. 计算布罗湖站流量（A和B的汇流）
    # A到罗湖站的延迟时段路径：A→三岔河口→文锦渡→罗湖站
    delay_dict = {
        '牛路岭到嘉积': (travel_time_period['牛路岭到嘉积']),
        '红岭到嘉积': travel_time_period['红岭到嘉积'],
    }

    # 1. A红岭到嘉积，
    Q_A_to_JJ = compute_flow(optimal_Q_A, delay_dict['牛路岭到嘉积'], T)
    # 2.牛路岭水库到布吉河口的流量
    Q_B_to_JJ = compute_flow(optimal_Q_B, delay_dict['红岭到嘉积'], T)
    Q_JJ = Q_A_to_JJ + Q_B_to_JJ
    # 嘉积水位
    H_JJ = (Q_JJ + 12955) / 1552.8
    # ---------------------- 目标函数：下游最大洪峰流量 ----------------------
    # 嘉积大坝的洪峰流量（过程最大值）
    peak_JJ = np.max(Q_JJ)
    # 下游整体最大洪峰（取所有水闸的最大值）
    total_peak_q = np.max([peak_JJ])
    total_peak_h = np.max([H_JJ])

    # 整体流量波动最小
    fluctuation_Q = 0
    for t in range(T - 1):
        fluctuation_Q += optimal_Q_A[t + 1] - optimal_Q_A[t] + optimal_Q_B[t + 1] - optimal_Q_B[t]

    total_v = 0
    for t in range(T):
        total_v += (
                (V_A[t + 1] - reservoir_params['A']['V_initial']) ** 2
                + (V_B[t + 1] - reservoir_params['B']['V_initial']) ** 2
        )

    total_b_v = 0
    for t in range(T):
        total_b_v += V_B[t + 1]

    # 提取最优解的索引
    optimal_index = np.argmin(res.F)

    # 从保存的流量过程中提取最优解的流量过程
    optimal_flow_process = problem.flow_processes[optimal_index]

    # 创建包含所有水闸流量的DataFrame
    flow_process_df = (pd.DataFrame({
        '时段': inflow_data['时间'],
        '红岭水库出库流量(m³/s)': optimal_Q_A,
        '牛路岭水库出库流量(m³/s)': optimal_Q_B,
        '': '',
        '红岭水库水位(m)': H_A[1:],
        '红岭水库库容(万m³)': V_A[1:],
        '牛路岭水库水位(m)': H_B[1:],
        '牛路岭水库库容(万m³)': V_B[1:],
        ' ': '',
        '嘉积流量(m³/s)': Q_JJ,
        '嘉积水位(m)': H_JJ,

    }))

    excel_name = '流量水位过程_海南5.xlsx'
    flow_process_df.to_excel(excel_name, index=False)

    # 输出提示信息
    print("\n调度过程已保存至" + excel_name + '。')

    end_time = time.perf_counter()  # 记录结束时间
    execution_time = end_time - start_time
    print(f"程序运行时间：{execution_time:.6f} 秒")

